#!/usr/bin/env python3
###
# @Author: Logan.Li
# @Gitee: https://gitee.com/attacker
# @email: admin@attacker.club
# @Date: 2025-03-14 10:00:30
# @LastEditTime: 2025-03-14 10:00:35
# @Description: AWS EC2监控模块
###

import logging
from typing import Dict, List
from botocore.exceptions import ClientError
from ..core import AWSMonitor

class EC2Monitor(AWSMonitor):
    """EC2监控类"""
    
    def __init__(self, config: Dict):
        super().__init__(config)
        self.namespace = "AWS/EC2"
        self.ec2 = self.session.client('ec2')
        
    def list_instances(self) -> List[Dict]:
        """获取所有EC2实例"""
        try:
            response = self.ec2.describe_instances()
            instances = []
            for reservation in response.get('Reservations', []):
                instances.extend(reservation.get('Instances', []))
            return instances
        except ClientError as e:
            logging.error(f"获取EC2实例列表失败: {e}")
            return []
            
    def get_instance_metrics(self, instance_id: str) -> Dict:
        """获取实例监控指标
        
        Args:
            instance_id: 实例ID
        """
        dimensions = [{
            'Name': 'InstanceId',
            'Value': instance_id
        }]
        
        metrics = {}
        metric_list = [
            ("CPUUtilization", "CPU使用率"),
            ("MemoryUtilization", "内存使用率"),  # 需要CloudWatch Agent
            ("DiskSpaceUtilization", "磁盘使用率"),  # 需要CloudWatch Agent
            ("NetworkIn", "网络入流量"),
            ("NetworkOut", "网络出流量")
        ]
        
        for metric_name, metric_desc in metric_list:
            data = self.get_metric_data(
                namespace=self.namespace,
                metric_name=metric_name,
                dimensions=dimensions
            )
            metrics[metric_desc] = data
            
        return metrics
        
    def monitor_all_instances(self) -> Dict:
        """监控所有EC2实例"""
        instances = self.list_instances()
        abnormal_instances = []
        
        for instance in instances:
            instance_id = instance['InstanceId']
            instance_name = ''
            # 获取实例名称标签
            for tag in instance.get('Tags', []):
                if tag['Key'] == 'Name':
                    instance_name = tag['Value']
                    break
            if not instance_name:
                instance_name = instance_id
                
            # 只监控运行中的实例
            if instance['State']['Name'] != 'running':
                continue
                
            # 获取实例指标
            metrics = self.get_instance_metrics(instance_id)
            
            # 检查异常指标
            for metric_name, data in metrics.items():
                if not data:
                    continue
                    
                # 获取最新数据点
                values = data.get('MetricDataResults', [])
                if not values or not values[0].get('Values'):
                    continue
                    
                value = values[0]['Values'][-1]
                
                # 根据不同指标判断异常
                is_abnormal = False
                if metric_name == "CPU使用率" and value > 80:
                    is_abnormal = True
                elif metric_name == "内存使用率" and value > 85:
                    is_abnormal = True
                elif metric_name == "磁盘使用率" and value > 90:
                    is_abnormal = True
                
                if is_abnormal:
                    abnormal_instances.append({
                        'id': instance_id,
                        'name': instance_name,
                        'metric': metric_name,
                        'value': value,
                        'time': data['MetricDataResults'][0]['Timestamps'][-1]
                    })
        
        # 聚合监控结果
        result = {
            'total_instances': len(instances),
            'running_instances': len([i for i in instances if i['State']['Name'] == 'running']),
            'abnormal_count': len(abnormal_instances),
            'top_abnormal': sorted(
                abnormal_instances,
                key=lambda x: x['value'],
                reverse=True
            )[:5]  # 只返回最异常的前5个实例
        }
        
        return result
